import warnings
warnings.filterwarnings('ignore')
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf
tf.compat.v1.logging.set_verbosity(40)
import os
import numpy as np
from scipy import signal
from scipy.io import wavfile
from sklearn.model_selection import train_test_split
from tensorflow.keras import optimizers, losses, activations, models
from tensorflow.keras.layers import Dense, Input, Flatten, MaxPooling2D, BatchNormalization,Conv2D,ReLU
import matplotlib.pyplot as plt
import random

# 分帧，加窗，傅里叶变换，取对数(倒谱)
def log_specgram(audio, sample_rate, window_size=20,
                 step_size=10, eps=1e-10):
    nperseg = int(round(window_size * sample_rate / 1e3))
    noverlap = int(round(step_size * sample_rate / 1e3))
    freqs, times, spec = signal.spectrogram(audio,
                                            fs=sample_rate,
                                            window='hann',
                                            nperseg=nperseg,
                                            noverlap=noverlap,
                                            detrend=False)
    return freqs, times, np.log(spec.T.astype(np.float32) + eps)

# 随机切分(类似数据增强)
def chop_audio(samples, L=16000, num=20):
    for i in range(num):
        beg = np.random.randint(0, len(samples) - L)
        yield samples[beg: beg + L]

# 定义模型
def ConvCell(inp,filters,kernel_size):
    x = Conv2D(filters,kernel_size,padding='same')(inp)
    x = BatchNormalization()(x)
    x = ReLU()(x)
    return x

def model_cnn(input_shape):
    nclass = 3
    inp = Input(shape=input_shape)
    norm_inp = BatchNormalization()(inp)

    img_1 = ConvCell(norm_inp,8,2)
    img_1 = ConvCell(img_1,8,2)
    img_1 = MaxPooling2D(pool_size=(2, 2))(img_1)

    img_1 = ConvCell(img_1,16,3)
    img_1 = ConvCell(img_1,16,3)
    img_1 = MaxPooling2D(pool_size=(2, 2))(img_1)

    img_1 = ConvCell(img_1, 32, 3)
    img_1 = ConvCell(img_1, 32, 3)
    img_1 = MaxPooling2D(pool_size=(2, 2))(img_1)

    img_1 = Flatten()(img_1)

    dense_1 = Dense(512, activation=activations.relu)(img_1)
    dense_1 = Dense(128, activation=activations.relu)(dense_1)
    dense_1 = Dense(nclass, activation=activations.softmax)(dense_1)

    model = models.Model(inputs=inp, outputs=dense_1)

    return model

def readvoice(path):
    y_data = []
    x_data = []
    for i, label in enumerate(os.listdir(path)):
        label_path = os.path.join(path, label)
        for voice_name in os.listdir(label_path):
            sample_rate, samples = wavfile.read(os.path.join(label_path, voice_name))
            samples = samples[:, 0]
            if len(samples) > 16000:
                n_samples = chop_audio(samples)
            else:
                n_samples = [samples]
            for j, samples in enumerate(n_samples):
                # signal.resample重采样
                resampled = signal.resample(samples, int(new_sample_rate / sample_rate * samples.shape[0]))
                _, _, specgram = log_specgram(resampled, sample_rate=new_sample_rate)
                y_data.append(np.eye(3)[i])
                x_data.append(np.expand_dims(specgram,2))

    return np.array(x_data),np.array(y_data)

if __name__ == "__main__":
    L = 16000
    train_path = r'E:\DATA\direction_data_3rd\train'
    test_path = r'E:\DATA\direction_data_3rd\test'

    new_sample_rate = 8000
    x_train,y_train = readvoice(train_path)
    x_test,y_test = readvoice(test_path)

    model = model_cnn(x_train.shape[1:])

    opt = optimizers.Adam()

    model.compile(optimizer=opt, loss=losses.categorical_crossentropy, metrics=['accuracy'])
    model.summary()

    x_train, x_valid, y_train, y_valid = train_test_split(x_train, y_train, test_size=0.1, random_state=2017)

    model.fit(x_train, y_train, batch_size=16, validation_data=(x_valid, y_valid), epochs=5, shuffle=True, verbose=2)

    score = model.evaluate(x_test, y_test, verbose=2)
    print(f'测试集损失值:{score[0]:.3f}')
    print(f'测试集准确率:{score[1]:.3f}')

    plt.figure(figsize=(12,12))
    for i in range(9):
        plt.subplot(3,3,i+1)
        r = random.randint(0,len(x_test)-1)
        plt.imshow(np.transpose(x_test[r],(1,0,2)))
        plt.ylim([0,81])
    plt.show()